# -*- coding: utf-8 -*-
from __future__ import print_function

import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim

from ddl_platform.ddlib import job

# for yolov3
import utils.datasets as data
from model.yolov3 import Yolov3
from model.loss.yolo_loss import YoloV3Loss

import os
#import models
import torchvision.models as models

class YOLOV3Job(job.Job):
    def build_dataset(self):
        config = self.config()['dataset']
        data_dir = config['data_dir']

        ## Data loading code
        traindir = os.path.join(data_dir, 'train')
        testdir = os.path.join(data_dir, 'val')

        trainset = data.VocDataset(anno_file_type="train", img_size=448)
        testset = data.VocDataset(anno_file_type="test")

        return trainset, testset


    def build_model(self):
        return Yolov3()

    def build_optimizer(self):
        model = self.model()
        config = self.config()['optimizer']
        optimizer = optim.SGD(model.parameters(), 
                lr=config['lr'],
                momentum=config['momentum'])
        return optimizer

    def build_criterion(self):
        criterion = YoloV3Loss(anchors=self.config()['loss']['anchors'], 
                strides=self.config()['loss']['strides'],
                iou_threshold_loss=self.config()['loss']['iou_threshold_loss'])
        return criterion

    def cal_eval_performance(self, batch_outputs, batch_inputs):

        with torch.no_grad():
            p, p_d = batch_outputs
            _, label_sbbox, label_mbbox, label_lbbox, sbboxes, mbbxes, lbboxes = batch_inputs
            criterion = YoloV3Loss(anchors=self.config()['loss']['anchors'], 
                    strides=self.config()['loss']['iou_threshold_loss'],
                    iou_threshold_loss=self.config()['loss']['iou_threshold_loss'])
            loss, loss_giou, loss_conf, loss_cls = self.criterion()(p, p_d, label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes)
            return loss_giou

    def training_step(self, batch, model):

        imgs, label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes = batch

        #imgs = imgs.cuda()
        #label_sbbox = label_sbbox.cuda()
        #label_mbbox = label_mbbox.cuda()
        #label_lbbox = label_lbbox.cuda()
        #sbboxes = sbboxes.cuda()
        #mbboxes = mbboxes.cuda()
        #lbboxes = lbboxes.cuda()

        p, p_d = model(imgs)

        loss, loss_giou, loss_conf, loss_cls = self.criterion()(p, p_d, label_sbbox, label_mbbox, label_lbbox, sbboxes, mbboxes, lbboxes)

        return loss, (p, p_d)


#def train():
#    conf_yaml = 'mnist.yaml'
#    job = MnistJob(conf_yaml)
#    t = trainer.Trainer(job)
#    t.fit()
#
#if __name__ == '__main__':
#    train()
